from crewai import Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task
from crewai_tools import DirectoryReadTool, FileReadTool

from mp_writer.myllms import MyLLM

# Uncomment the following line to use an example of a custom tool
from mp_writer.tools.csdn import CSDNSearchTool, CSDNViewArticleTool
from crewai_tools import SerperDevTool
from langchain.agents import load_tools

# Check our tools documentations for more information on how to use them
# from crewai_tools import SerperDevTool
llm = MyLLM(model_name="glm-4").get_llm()
llm2 = MyLLM(model_name="gpt-4").get_llm()
search_tool = SerperDevTool()
human_tools = load_tools(["human"])
dic_tool = DirectoryReadTool(directory='/home/jialin/Documents/AI/学习笔记/资料整理/crewAI/')
file_read_tool = FileReadTool()


@CrewBase
class MpWriterCrew():
	"""MpWriter crew"""
	agents_config = 'config/agents.yaml'
	tasks_config = 'config/tasks.yaml'

	@agent
	def researcher(self) -> Agent:
		return Agent(
			config=self.agents_config['researcher'],
			# tools=[CSDNSearchTool(), CSDNViewArticleTool()], 
			# tools=[search_tool], 
			tools = [dic_tool, file_read_tool],
			verbose=True,
			llm=llm,
			function_calling_llm=llm2,
			allow_delegation=False
		)

	@agent
	def reporting_analyst(self) -> Agent:
		return Agent(
			config=self.agents_config['reporting_analyst'],
			tool=human_tools,
			verbose=True,
			llm=llm,
			function_calling_llm=llm2,
		)

	@task
	def research_task(self) -> Task:
		return Task(
			config=self.tasks_config['research_task'],
			agent=self.researcher()
		)

	@task
	def reporting_task(self) -> Task:
		return Task(
			config=self.tasks_config['reporting_task'],
			agent=self.reporting_analyst(),
			output_file='report.md'
		)

	@crew
	def crew(self) -> Crew:
		"""Creates the MpWriter crew"""
		return Crew(
			agents=self.agents, # Automatically created by the @agent decorator
			tasks=self.tasks, # Automatically created by the @task decorator
			process=Process.sequential,
			verbose=2,
			# process=Process.hierarchical, # In case you wanna use that instead https://docs.crewai.com/how-to/Hierarchical/
		)